#Need to uncomment and run this when the container is first instantiated
#!apt-get install python3-opencv
import cv2 as cv
from PIL import Image
import sys
import os
import urllib
import tensorflow.contrib.tensorrt as trt
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import tensorflow as tf
import numpy as np
import time
from tf_trt_models.detection import download_detection_model, build_detection_graph
%matplotlib inline
# https://github.com/yeephycho/tensorflow-face-detection
FROZEN_GRAPH_NAME = 'data/frozen_inference_graph_face.pb'
!wget https://github.com/yeephycho/tensorflow-face-detection/blob/master/model/frozen_inference_graph_face.pb?raw=true -O {FROZEN_GRAPH_NAME}
cv.imwrite('data/my_sample_webcam_pic.png',img)
# # go warriors!
# IMAGE_PATH = 'data/warriors.jpg'
# !wget 'https://cdn.vox-cdn.com/thumbor/rC0mlBATZdoDW1tEa44P6431sGc=/0x0:3683x2455/1200x800/filters:focal(1623x234:2211x822)/cdn.vox-cdn.com/uploads/chorus_image/image/63273148/usa_today_12005182.0.jpg' -O {IMAGE_PATH}
output_dir=''
frozen_graph = tf.GraphDef()
with open(os.path.join(output_dir, FROZEN_GRAPH_NAME), 'rb') as f:
frozen_graph.ParseFromString(f.read())
# https://github.com/NVIDIA-AI-IOT/tf_trt_models/blob/master/tf_trt_models/detection.py
INPUT_NAME='image_tensor'
BOXES_NAME='detection_boxes'
CLASSES_NAME='detection_classes'
SCORES_NAME='detection_scores'
MASKS_NAME='detection_masks'
NUM_DETECTIONS_NAME='num_detections'
input_names = [INPUT_NAME]
output_names = [BOXES_NAME, CLASSES_NAME, SCORES_NAME, NUM_DETECTIONS_NAME]
trt_graph = trt.create_inference_graph(
input_graph_def=frozen_graph,
outputs=output_names,
max_batch_size=1,
max_workspace_size_bytes=1 << 25,
precision_mode='FP16',
minimum_segment_size=50
)
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
tf_sess = tf.Session(config=tf_config)
# use this if you want to try on the optimized TensorRT graph
# Note that this will take a while
# tf.import_graph_def(trt_graph, name='')
# use this if you want to try directly on the frozen TF graph
# this is much faster
tf.import_graph_def(frozen_graph, name='')
tf_input = tf_sess.graph.get_tensor_by_name(input_names[0] + ':0')
tf_scores = tf_sess.graph.get_tensor_by_name('detection_scores:0')
tf_boxes = tf_sess.graph.get_tensor_by_name('detection_boxes:0')
tf_classes = tf_sess.graph.get_tensor_by_name('detection_classes:0')
tf_num_detections = tf_sess.graph.get_tensor_by_name('num_detections:0')
# suppress boxes that are below the threshold..
DETECTION_THRESHOLD = 0.5
def get_images(cap):
_ret, img = cap.read()
if not _ret:
print('No image')
return
img_color = cv.cvtColor(img, cv.COLOR_BGR2RGB)
#print("The image as obtained from the webcam")
#plt.imshow(img_color)
#plt.show()
image = Image.fromarray(img_color)
image_resized = np.array(image.resize((300, 300)))
#image_resized = np.resize(image, (300,300))
#print(image_resized.shape)
#print("The image that has been resized to run inference")
#plt.imshow(image_resized)
#plt.show()
image = np.array(image)
run_inference(image_resized, image)
def run_inference(image_resized, image):
scores, boxes, classes, num_detections = tf_sess.run(
[tf_scores, tf_boxes, tf_classes, tf_num_detections],
feed_dict={ tf_input: image_resized[None, ...] })
boxes = boxes[0] # index by 0 to remove batch dimension
scores = scores[0]
classes = classes[0]
num_detections = num_detections[0]
display_detected_images(image, boxes, scores, num_detections, classes)
def display_detected_images(image, boxes, scores, num_detections, classes):
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.imshow(image)
# plot boxes exceeding score threshold
for i in range(int(num_detections)):
if scores[i] < DETECTION_THRESHOLD:
continue
# scale box to image coordinates
box = boxes[i] * np.array([image.shape[0], image.shape[1], image.shape[0], image.shape[1]])
# display rectangle
patch = patches.Rectangle((box[1], box[0]), box[3] - box[1], box[2] - box[0], color='g', alpha=0.3)
ax.add_patch(patch)
# display class index and score
plt.text(x=box[1] + 10, y=box[2] - 10, s='%d (%0.2f) ' % (classes[i], scores[i]), color='w')
print("Image with face being detected with the box")
plt.show()
cap = cv.VideoCapture(1)
for i in range(30):
get_images(cap)
cap.release()
num_samples = 50
t0 = time.time()
for i in range(num_samples):
scores, boxes, classes, num_detections = tf_sess.run([tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={
tf_input: image_resized[None, ...]
})
t1 = time.time()
print('Average runtime: %f seconds' % (float(t1 - t0) / num_samples))
tf_sess.close()